package com.vlefen.springcloud.grpc.consul;

import com.ecwid.consul.v1.agent.model.NewCheck;
import com.google.common.collect.Sets;
import com.vlefen.springcloud.grpc.Registry;
import com.vlefen.springcloud.grpc.config.GRpcProperties;
import io.grpc.ServerServiceDefinition;
import org.springframework.cloud.client.serviceregistry.AutoServiceRegistrationProperties;
import org.springframework.cloud.consul.discovery.ConsulDiscoveryProperties;
import org.springframework.cloud.consul.serviceregistry.ConsulAutoRegistration;
import org.springframework.cloud.consul.serviceregistry.ConsulServiceRegistry;
import org.springframework.context.ApplicationContext;
import org.springframework.util.CollectionUtils;

import java.util.Set;

public class ConsulRegistry implements Registry {

    private final ConsulServiceRegistry consulServiceRegistry;
    private final ApplicationContext context;
    private final ConsulDiscoveryProperties discoveryProperties;
    private final GRpcProperties gRpcProperties;
    private final Set<ServerServiceDefinition> serverServiceDefinitions;
    private final Set<ConsulAutoRegistration> registrations = Sets.newConcurrentHashSet();


    public ConsulRegistry(Set<ServerServiceDefinition> serverServiceDefinitions,
                          ConsulServiceRegistry consulServiceRegistry,
                          ConsulDiscoveryProperties discoveryProperties,
                          GRpcProperties gRpcProperties,
                          ApplicationContext context) {
        this.serverServiceDefinitions = serverServiceDefinitions;
        this.consulServiceRegistry = consulServiceRegistry;
        this.context = context;
        this.discoveryProperties = discoveryProperties;
        this.gRpcProperties = gRpcProperties;
    }


    public void register() {
        if (!CollectionUtils.isEmpty(serverServiceDefinitions)) {
            serverServiceDefinitions.stream().forEach(this::register);
        }
    }

    private void register(ServerServiceDefinition serverServiceDefinition) {
        final String serviceName = serverServiceDefinition.getServiceDescriptor().getName();
        ConsulService service = new ConsulService();
        String appName = ConsulAutoRegistration.getAppName(discoveryProperties, context.getEnvironment());
        service.setId(appName + "-" + serviceName);
        if (!discoveryProperties.isPreferAgentAddress()) {
            service.setAddress(discoveryProperties.getHostname());
        }

        service.setName(serviceName);
        service.setTags(ConsulAutoRegistration.createTags(discoveryProperties));
        service.setAddress(gRpcProperties.getServer().getAddress());
        service.setPort(gRpcProperties.getServer().getPort());


        NewCheck newCheck = new NewCheck();
        newCheck.setInterval(gRpcProperties.getServer().getHealthCheckInterval());
        newCheck.setGrpc(gRpcProperties.getServer().getAddress() + ":" + gRpcProperties.getServer().getPort() + "/" + serviceName);
        newCheck.setGrpcUseTLS(gRpcProperties.getServer().getSsl().isEnabled());
        newCheck.setName("GRpc service " + service.getName() + " health check");
        newCheck.setId(service.getId() + "_health_check");
        newCheck.setServiceId(service.getId());
        service.check = newCheck;
        AutoServiceRegistrationProperties autoServiceRegistrationProperties = new AutoServiceRegistrationProperties();
        autoServiceRegistrationProperties.setEnabled(true);
        autoServiceRegistrationProperties.setFailFast(gRpcProperties.getServer().isFailFast());
        ConsulAutoRegistration registration = new ConsulAutoRegistration(service, autoServiceRegistrationProperties, discoveryProperties, context, null);
        consulServiceRegistry.register(registration);
        registrations.add(registration);
    }

    public void deregister() {

        if (null != consulServiceRegistry) {
            try {
                registrations.stream().forEach(consulServiceRegistry::deregister);
            } catch (Throwable throwable) {

            }
        }
    }
}
